{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "The first experiment!\n",
    "This code is used to predict truncated Galois LFSRs with known parameters.\n",
    "Suppose that the modulus and the state transition matrix are known, given the first Ni (i=1,2,...,n) high-order truncated digits of the i-th sequence output by the i-th register, we need to recover the initial state.\n",
    "The library used is the Sagemath.\n",
    "'''\n",
    "\n",
    "# Choose n1 registers from n registers and set the length of each sequence output by register. The total number of output digits is equal to d.\n",
    "def distribution(n, n1, d):\n",
    "    N = vector(ZZ, n1)\n",
    "    D = vector(ZZ, n)\n",
    "    C = Combinations(range(n), n1)\n",
    "    N = C.list()[randint(0, C.cardinality() - 1)]\n",
    "    for i in N:\n",
    "        D[i] = 1\n",
    "    for i in range(n1, d):\n",
    "        k = randint(0, n1 - 1)\n",
    "        D[N[k]] = D[N[k]] + 1\n",
    "    return N, D\n",
    "\n",
    "# Recover the initial state by the high-order truncated digits\n",
    "def recover_the_initial_state_by_high_order_truncated_digits(m, n, d, n1, beta):\n",
    "    C = random_matrix(ZZ, n, n, x = 0, y = m)\n",
    "    A = random_matrix(ZZ, d, n, x = 0, y = m)\n",
    "    B = Matrix(ZZ, d, n)\n",
    "    Y = Matrix(ZZ, d, n)\n",
    "    N, D = distribution(n, n1, d)\n",
    "    n2 = n - n1\n",
    "    N2 = vector(ZZ, n2)\n",
    "    a = Matrix(ZZ, d + 2 * n2 + 1, n)\n",
    "    \n",
    "    i2 = 0\n",
    "    for i in range(0, n):\n",
    "        if i in N:\n",
    "            0\n",
    "        else:\n",
    "            N2[i2] = i\n",
    "            i2 = i2 + 1\n",
    "            \n",
    "    for i in range(1, d):\n",
    "        A[i] = (A[i - 1] * C) % m\n",
    "    for i in range(0, n):\n",
    "        for j in range(0, D[i]):\n",
    "            Y[j, i] = A[j, i] >> beta\n",
    "    \n",
    "    L = Matrix(ZZ, d + 2 * n2 + 1, d + 2 * n2 + 1)\n",
    "    L[0, 0] = 2^(beta - 1)\n",
    "    for i in range(1, n + 1):\n",
    "        L[0, i] = 2^(beta - 1)\n",
    "        L[i, i] = 1\n",
    "    for i in range(n + 1, n + n2 + 1):\n",
    "        L[0, i] = 2^(beta - 1)\n",
    "    col = n + n2\n",
    "    \n",
    "    for i in N:\n",
    "        CC = identity_matrix(ZZ, n)\n",
    "        for j in range(1, D[i]):\n",
    "            CC = (CC * C) % m\n",
    "            B[j] = (2^beta * (Y[0] * CC - Y[j])) % m\n",
    "            L[0, col + j] = -B[j, i] + 2^(beta - 1)\n",
    "            for k in range(0, n2):\n",
    "                L[n + k + 1, col + j] = 2^beta * CC[N2[k], i]\n",
    "            L[col + j, col + j] = m\n",
    "            for k in range(0, n):\n",
    "                L[k + 1, col + j] = CC[k, i]\n",
    "        col = col + D[i] - 1\n",
    "        \n",
    "    if beta >= alpha:\n",
    "        for i in range(n + 1, n + n2 + 1):\n",
    "            L[i, i] = 2^(beta - alpha)\n",
    "    else:\n",
    "        for i in range(0, d + 2 * n2 + 1):\n",
    "            for j in range(0, d + 2 * n2 + 1):\n",
    "                L[i, j] = L[i, j] * 2^(alpha - beta)\n",
    "        for i in range(n + 1, n + n2 + 1):\n",
    "            L[i, i] = 1\n",
    "            \n",
    "    reduced_L = L.BKZ(block_size = 20)\n",
    "\n",
    "    flag = 0\n",
    "    if beta >= alpha:\n",
    "        for i in range(0, d + 2 * n2 + 1):\n",
    "            if reduced_L[i, 0] == 2^(beta - 1):\n",
    "                for j in range(1, n + 1):\n",
    "                    reduced_L[i, j] = -reduced_L[i, j] + 2^(beta - 1)\n",
    "                for j in range(n + 1, n + n2 + 1):\n",
    "                    reduced_L[i, j] = (-reduced_L[i, j] + 2^(beta - 1)) * 2^alpha\n",
    "                for j in range(0, n2):\n",
    "                    a[flag, N2[j]] = reduced_L[i, N2[j] + 1] + reduced_L[i, n + j + 1]\n",
    "                for j in range(0, n1):\n",
    "                    a[flag, N[j]] = reduced_L[i, N[j] + 1] + Y[0, N[j]] * 2^beta\n",
    "                flag = flag + 1\n",
    "            if reduced_L[i, 0] == -2^(beta - 1):\n",
    "                for j in range(1, n + 1):\n",
    "                    reduced_L[i, j] = reduced_L[i, j] + 2^(beta - 1)\n",
    "                for j in range(n + 1, n + n2 + 1):\n",
    "                    reduced_L[i, j] = (reduced_L[i, j] + 2^(beta - 1)) * 2^alpha\n",
    "                for j in range(0, n2):\n",
    "                    a[flag, N2[j]] = reduced_L[i, N2[j] + 1] + reduced_L[i, n + j + 1]\n",
    "                for j in range(0, n1):\n",
    "                    a[flag, N[j]] = reduced_L[i, N[j] + 1] + Y[0, N[j]] * 2^beta\n",
    "                flag = flag + 1\n",
    "\n",
    "    if beta < alpha:\n",
    "        for i in range(0, d + 2 * n2 + 1):\n",
    "            if reduced_L[i, 0] == 2^(alpha - 1):\n",
    "                for j in range(1, n + 1):\n",
    "                    reduced_L[i, j] = (-reduced_L[i, j] + 2^(alpha - 1)) / 2^(alpha - beta)\n",
    "                for j in range(n + 1, n + n2 + 1):\n",
    "                    reduced_L[i, j] = -reduced_L[i, j] + 2^(alpha - 1)\n",
    "                for j in range(0, n2):\n",
    "                    a[flag, N2[j]] = reduced_L[i, N2[j] + 1] + reduced_L[i, n + j + 1] * 2^beta\n",
    "                for j in range(0, n1):\n",
    "                    a[flag, N[j]] = reduced_L[i, N[j] + 1] + Y[0, N[j]] * 2^beta\n",
    "                flag = flag + 1\n",
    "            if reduced_L[i, 0] == -2^(alpha - 1):\n",
    "                for j in range(1, n + 1):\n",
    "                    reduced_L[i, j] = (reduced_L[i, j] + 2^(alpha - 1)) / 2^(alpha - beta)\n",
    "                for j in range(n + 1, n + n2 + 1):\n",
    "                    reduced_L[i, j] = reduced_L[i, j] + 2^(alpha - 1)\n",
    "                for j in range(0, n2):\n",
    "                    a[flag, N2[j]] = reduced_L[i, N2[j] + 1] + reduced_L[i, n + j + 1] * 2^beta\n",
    "                for j in range(0, n1):\n",
    "                    a[flag, N[j]] = reduced_L[i, N[j] + 1] + Y[0, N[j]] * 2^beta\n",
    "                flag = flag + 1\n",
    "                \n",
    "    return A[0], a\n",
    "\n",
    "# Determine whether the initial state is recovered \n",
    "def judge(m, n, d, n1, beta):\n",
    "    A, a = recover_the_initial_state_by_high_order_truncated_digits(m, n, d, n1, beta)\n",
    "    if A == a[0]:\n",
    "        return 1\n",
    "    else:\n",
    "        return 0\n",
    "\n",
    "\n",
    "n = 16\n",
    "m = 2^31 - 1\n",
    "\n",
    "alpha = 17\n",
    "n1 = 4\n",
    "\n",
    "n2 = n - n1\n",
    "beta = 31 - alpha\n",
    "s = 100\n",
    "for d in range(31, 20, -1):\n",
    "    tol = 0\n",
    "    for i in range(0, s):\n",
    "        jud = judge(m, n, d, n1, beta)\n",
    "        if jud:\n",
    "            tol = tol + 1\n",
    "    print(d, tol)\n",
    "\n",
    "\n",
    "\n",
    "'''\n",
    "The second experiment!\n",
    "This code is used to predict truncated Galois LFSRs with known parameters.\n",
    "Suppose that the modulus and the state transition matrix are known, given the first Ni (i=1,2,...,n) low-order truncated digits of the i-th sequence output by the i-th register, we need to recover the initial state.\n",
    "The library used is the Sagemath.\n",
    "'''\n",
    "\n",
    "# Choose n1 registers from n registers and set the length of each sequence output by register. The total number of output digits is equal to d.\n",
    "def distribution(n, n1, d):\n",
    "    N = vector(ZZ, n1)\n",
    "    D = vector(ZZ, n)\n",
    "    C = Combinations(range(n), n1)\n",
    "    N = C.list()[randint(0, C.cardinality() - 1)]\n",
    "    for i in N:\n",
    "        D[i] = 1\n",
    "    for i in range(n1, d):\n",
    "        k = randint(0, n1 - 1)\n",
    "        D[N[k]] = D[N[k]] + 1\n",
    "    return N, D\n",
    "\n",
    "# Recover the initial state by the low-order truncated digits\n",
    "def recover_the_initial_state_by_low_order_truncated_digits(m, n, d, n1, alpha, beta):\n",
    "    C = random_matrix(ZZ, n, n, x = 0, y = m)\n",
    "    A = random_matrix(ZZ, d, n, x = 0, y = m)\n",
    "    B = Matrix(ZZ, d, n)\n",
    "    Z = Matrix(ZZ, d, n)\n",
    "    N, D = distribution(n, n1, d)\n",
    "    n2 = n - n1\n",
    "    N2 = vector(ZZ, n2)\n",
    "    dim = d + 2 * n2 + 1\n",
    "    a = Matrix(ZZ, 2 * dim, n)\n",
    "    \n",
    "    i2 = 0\n",
    "    for i in range(0, n):\n",
    "        if i in N:\n",
    "            0\n",
    "        else:\n",
    "            N2[i2] = i\n",
    "            i2 = i2 + 1\n",
    "            \n",
    "    for i in range(1, d):\n",
    "        A[i] = (A[i - 1] * C) % m\n",
    "    for i in range(0, n):\n",
    "        for j in range(0, D[i]):\n",
    "            Z[j, i] = A[j, i] % (2^beta)\n",
    "    \n",
    "    gamma = gcd(m, 2^beta)\n",
    "    m_new = m / gamma\n",
    "     \n",
    "    L = Matrix(ZZ, dim, dim)\n",
    "    L[0, 0] = 2^(alpha - 1)\n",
    "    for i in range(1, n + 1):\n",
    "        L[0, i] = 2^(alpha - 1)\n",
    "        L[i, i] = 1\n",
    "    for i in range(n + 1, n + n2 + 1):\n",
    "        L[0, i] = 2^(alpha - 1)\n",
    "    col = n + n2\n",
    "    \n",
    "    for i in N:\n",
    "        CC = identity_matrix(ZZ, n)\n",
    "        for j in range(1, D[i]):\n",
    "            CC = (CC * C) % m\n",
    "            B[j] = (2^(-beta) * (Z[0] * CC - Z[j])) % m_new\n",
    "            L[0, col + j] = -B[j, i] + 2^(alpha - 1)\n",
    "            for k in range(0, n2):\n",
    "                L[n + k + 1, col + j] = (2^(-beta) * CC[N2[k], i]) % m_new\n",
    "            L[col + j, col + j] = m_new\n",
    "            for k in range(0, n):\n",
    "                L[k + 1, col + j] = CC[k, i] % m_new\n",
    "        col = col + D[i] - 1\n",
    "        \n",
    "    if alpha >= beta:\n",
    "        for i in range(n + 1, n + n2 + 1):\n",
    "            L[i, i] = 2^(alpha - beta)\n",
    "    else:\n",
    "        for i in range(0, dim):\n",
    "            for j in range(0, dim):\n",
    "                L[i, j] = L[i, j] * 2^(beta - alpha)\n",
    "        for i in range(n + 1, n + n2 + 1):\n",
    "            L[i, i] = 1 \n",
    "\n",
    "    reduced_L = L.BKZ(block_size = 20)\n",
    "\n",
    "    flag = 0\n",
    "    if alpha >= beta:\n",
    "        for i in range(0, dim):\n",
    "            if reduced_L[i, 0] == 2^(alpha - 1):                \n",
    "                for j in range(1, n + 1):\n",
    "                    reduced_L[i, j] = -reduced_L[i, j] + 2^(alpha - 1)\n",
    "                for j in range(n + 1, n + n2 + 1):\n",
    "                    reduced_L[i, j] = (-reduced_L[i, j] + 2^(alpha - 1)) / 2^(alpha - beta)\n",
    "                for j in range(0, n2):\n",
    "                    a[flag, N2[j]] = reduced_L[i, N2[j] + 1] * 2^beta + reduced_L[i, n + j + 1]\n",
    "                for j in range(0, n1):\n",
    "                    a[flag, N[j]] = reduced_L[i, N[j] + 1] * 2^beta + Z[0, N[j]]\n",
    "                flag = flag + 1\n",
    "                for j in range(0, n2):\n",
    "                    if reduced_L[i, N2[j] + 1] < 2^(alpha - 1):\n",
    "                        a[flag, N2[j]] = (reduced_L[i, N2[j] + 1] + 2^(alpha - 1)) * 2^beta + reduced_L[i, n + j + 1] - 1\n",
    "                    if reduced_L[i, N2[j] + 1] > 2^(alpha - 1):\n",
    "                        a[flag, N2[j]] = (reduced_L[i, N2[j] + 1] - 2^(alpha - 1)) * 2^beta + reduced_L[i, n + j + 1] + 1\n",
    "                flag = flag + 1\n",
    "                              \n",
    "            if reduced_L[i, 0] == -2^(alpha - 1):\n",
    "                for j in range(1, n + 1):\n",
    "                    reduced_L[i, j] = reduced_L[i, j] + 2^(alpha - 1)\n",
    "                for j in range(n + 1, n + n2 + 1):\n",
    "                    reduced_L[i, j] = (reduced_L[i, j] + 2^(alpha - 1)) / 2^(alpha - beta)\n",
    "                for j in range(0, n2):\n",
    "                    a[flag, N2[j]] = reduced_L[i, N2[j] + 1] * 2^beta + reduced_L[i, n + j + 1]\n",
    "                for j in range(0, n1):\n",
    "                    a[flag, N[j]] = reduced_L[i, N[j] + 1] * 2^beta + Z[0, N[j]]\n",
    "                flag = flag + 1\n",
    "                for j in range(0, n2):\n",
    "                    if reduced_L[i, N2[j] + 1] < 2^(alpha - 1):\n",
    "                        a[flag, N2[j]] = (reduced_L[i, N2[j] + 1] + 2^(alpha - 1)) * 2^beta + reduced_L[i, n + j + 1] - 1\n",
    "                    if reduced_L[i, N2[j] + 1] > 2^(alpha - 1):\n",
    "                        a[flag, N2[j]] = (reduced_L[i, N2[j] + 1] - 2^(alpha - 1)) * 2^beta + reduced_L[i, n + j + 1] + 1\n",
    "                flag = flag + 1\n",
    "                \n",
    "\n",
    "    if alpha < beta:\n",
    "        for i in range(0, dim):\n",
    "            if reduced_L[i, 0] == 2^(beta - 1):\n",
    "                for j in range(1, n + 1):\n",
    "                    reduced_L[i, j] = (-reduced_L[i, j] + 2^(beta - 1)) / 2^(beta - alpha)\n",
    "                for j in range(n + 1, n + n2 + 1):\n",
    "                    reduced_L[i, j] = -reduced_L[i, j] + 2^(beta - 1)\n",
    "                for j in range(0, n2):\n",
    "                    a[flag, N2[j]] = reduced_L[i, N2[j] + 1] * 2^beta + reduced_L[i, n + j + 1]\n",
    "                for j in range(0, n1):\n",
    "                    a[flag, N[j]] = reduced_L[i, N[j] + 1] * 2^beta + Z[0, N[j]]\n",
    "                flag = flag + 1\n",
    "                for j in range(0, n2):\n",
    "                    if reduced_L[i, N2[j] + 1] < 2^(alpha - 1):\n",
    "                        a[flag, N2[j]] = (reduced_L[i, N2[j] + 1] + 2^(alpha - 1)) * 2^beta + reduced_L[i, n + j + 1] - 1\n",
    "                    if reduced_L[i, N2[j] + 1] > 2^(alpha - 1):\n",
    "                        a[flag, N2[j]] = (reduced_L[i, N2[j] + 1] - 2^(alpha - 1)) * 2^beta + reduced_L[i, n + j + 1] + 1\n",
    "                flag = flag + 1\n",
    "                \n",
    "            if reduced_L[i, 0] == -2^(beta - 1):\n",
    "                for j in range(1, n + 1):\n",
    "                    reduced_L[i, j] = (reduced_L[i, j] + 2^(beta - 1)) / 2^(beta - alpha)\n",
    "                for j in range(n + 1, n + n2 + 1):\n",
    "                    reduced_L[i, j] = reduced_L[i, j] + 2^(beta - 1)\n",
    "                for j in range(0, n2):\n",
    "                    a[flag, N2[j]] = reduced_L[i, N2[j] + 1] * 2^beta + reduced_L[i, n + j + 1]\n",
    "                for j in range(0, n1):\n",
    "                    a[flag, N[j]] = reduced_L[i, N[j] + 1] * 2^beta + Z[0, N[j]]\n",
    "                flag = flag + 1\n",
    "                for j in range(0, n2):\n",
    "                    if reduced_L[i, N2[j] + 1] < 2^(alpha - 1):\n",
    "                        a[flag, N2[j]] = (reduced_L[i, N2[j] + 1] + 2^(alpha - 1)) * 2^beta + reduced_L[i, n + j + 1] - 1\n",
    "                    if reduced_L[i, N2[j] + 1] > 2^(alpha - 1):\n",
    "                        a[flag, N2[j]] = (reduced_L[i, N2[j] + 1] - 2^(alpha - 1)) * 2^beta + reduced_L[i, n + j + 1] + 1\n",
    "                flag = flag + 1\n",
    "                \n",
    "    return A[0], a\n",
    "\n",
    "# Determine whether the initial state is recovered \n",
    "def judge(m, n, d, n1, alpha, beta):\n",
    "    A, a = recover_the_initial_state_by_low_order_truncated_digits(m, n, d, n1, alpha, beta)\n",
    "    count = 0\n",
    "    for j in range(0, n):\n",
    "        if A[j] == a[0, j] or A[j] == a[1, j]:\n",
    "            count = count + 1\n",
    "    if count == n:\n",
    "        return 1\n",
    "    else:\n",
    "        return 0\n",
    "    \n",
    "\n",
    "n = 16\n",
    "m = 2^31 - 2\n",
    "\n",
    "beta = 17\n",
    "n1 = 4\n",
    "\n",
    "n2 = n - n1\n",
    "alpha = 31 - beta\n",
    "s = 100\n",
    "\n",
    "for d in range(32, 20, -1):\n",
    "    tol = 0\n",
    "    for i in range(0, s):\n",
    "        result = judge(m, n, d, n1, alpha, beta)\n",
    "        if result:\n",
    "            tol = tol + 1\n",
    "    print(d, tol)\n",
    "\n",
    "\n",
    "\n",
    "'''\n",
    "The third experiment!\n",
    "This code is used to predict truncated Galois LFSRs with unknown parameters.\n",
    "Suppose that the modulus and the state transition matrix are unknown, given the first Ni (Ni>r, i=1,2,...,n) low-order truncated digits of the i-th sequence output by the i-th register, we need to recover the modulus, the characteristic polynomial, the initial state, and the state transition matrix.\n",
    "The library used is the Sagemath.\n",
    "'''\n",
    "import time \n",
    "\n",
    "# Recover the modulus by the resultant    \n",
    "def recover_the_modulus(m, n, r, t, D, A, Y):\n",
    "    time_start = time.clock()\n",
    "    L = Matrix(ZZ, r, t + r)    \n",
    "    for i in range(0, r):\n",
    "        L[i, t + i] = 1\n",
    "        p = 0\n",
    "        for j in range(0, n):\n",
    "            for k in range(0, D[j] - r + 1):\n",
    "                L[i, p + k] = Y[i + k, j]\n",
    "            p = p + D[j] - r + 1      \n",
    "    reduced_L = L.BKZ(block_size = 20)\n",
    "    \n",
    "    count = vector(ZZ, r)\n",
    "    for i in range(0, r):\n",
    "        yy = 0\n",
    "        for j in range(0, n):\n",
    "            for k in range(0, D[j] - r + 1):\n",
    "                z = 0\n",
    "                for zz in range(0, r):\n",
    "                    z = z + reduced_L[i, t + zz] * A[k + zz, j]\n",
    "                if z == 0:\n",
    "                    yy = yy + 1\n",
    "                if z != 0:\n",
    "                    yy = 0\n",
    "                    break\n",
    "            if yy == 0:\n",
    "                break\n",
    "        if yy == t:\n",
    "            count[i] = count[i] + 1\n",
    "    time_end = time.clock()\n",
    "    print('Step 1: time of searching polynomials that annihilate the sequence: ', time_end - time_start, 's\\n')\n",
    "    \n",
    "    time_start = time.clock()\n",
    "    R.<x> = ZZ[]\n",
    "    f = []\n",
    "    fi = reduced_L[0, t] * x^0\n",
    "    for i in range(1, r):\n",
    "        fi = fi + x^i * reduced_L[0, i + t]\n",
    "    f.append(fi)\n",
    "    fi = reduced_L[1, t] * x^0\n",
    "    for i in range(1, r):\n",
    "        fi = fi + x^i * reduced_L[1, i + t]\n",
    "    f.append(fi)\n",
    "    recovered_m_n = f[0].resultant(f[1])\n",
    "    \n",
    "    flag = 2\n",
    "    while recovered_m_n != m^n:\n",
    "        if flag == r - t:\n",
    "            flag = -1\n",
    "            break\n",
    "        fi = reduced_L[flag, t] * x^0\n",
    "        for i in range(1, r):\n",
    "            fi = fi + x^i * reduced_L[flag, i + t]\n",
    "        f.append(fi)\n",
    "        for i in range(0, flag):\n",
    "            res = f[flag].resultant(f[i])\n",
    "            recovered_m_n = gcd(recovered_m_n, res)\n",
    "        flag = flag + 1\n",
    "    \n",
    "    if flag > 0:\n",
    "        recovered_m = recovered_m_n^(1/n)\n",
    "        print(\"Success to recover m by \", flag, \" polynomials!\")\n",
    "        print(\"m == \", recovered_m, \"\\n\")\n",
    "    else:\n",
    "        print(\"Fail to recover m!\")\n",
    "        print(\"The wrong m == \", recovered_m_n^(1/n), \"\\n\")\n",
    "    \n",
    "    time_end = time.clock()\n",
    "    print('Step 2: time of recovering the modulus: ', time_end - time_start, 's\\n')\n",
    "    return count, reduced_L, flag\n",
    "\n",
    "# Recover the coefficients by the Chinese Remainder Theorem\n",
    "def recover_the_coefficients(L, m, F, r, t):\n",
    "    time_start = time.clock()\n",
    "    R.<x> = PolynomialRing(Integers(m), implementation='NTL')\n",
    "    \n",
    "    for i in range(0, r):\n",
    "        mark = r - 1\n",
    "        while L[i, mark + t] == 0:\n",
    "            mark = mark - 1\n",
    "        for j in range(0, t + r):\n",
    "            L[i, j] = L[i, j] * ((L[i, mark + t]^(-1)) % m)\n",
    "    \n",
    "    f = []\n",
    "    fi = L[0, t]\n",
    "    for i in range(1, r):\n",
    "        fi = fi + x^i * L[0, i + t]\n",
    "    f.append(fi)\n",
    "    fi = L[1, t]\n",
    "    for i in range(1, r):\n",
    "        fi = fi + x^i * L[1, i + t]\n",
    "    f.append(fi)\n",
    "    g = f[0].gcd(f[1])\n",
    "    \n",
    "    flag = 2\n",
    "    while g != F:\n",
    "        if flag == r - t:\n",
    "            flag = -1\n",
    "            break\n",
    "        fi = L[flag, t]\n",
    "        for i in range(1, r):\n",
    "            fi = fi + x^i * L[flag, i + t]\n",
    "        f.append(fi)\n",
    "        for i in range(0, flag):\n",
    "            g = g.gcd(f[flag])\n",
    "        flag = flag + 1\n",
    "    \n",
    "    if flag > 0:\n",
    "        print(\"Success to recover f by \", flag, \" polynomials!\")\n",
    "        print(\"f == \", g, \"\\n\")\n",
    "    else:\n",
    "        print(\"Fail to recover f!\")\n",
    "        print(\"The correct f == \", F, \"\\n\")\n",
    "        print(\"The wrong f == \", g, \"\\n\")\n",
    "        \n",
    "    time_end = time.clock()\n",
    "    print('Step 3: time of recovering the coefficients: ', time_end - time_start, 's\\n')\n",
    "    return flag\n",
    "\n",
    "# Recover the initial state by Kannan's embedding technique\n",
    "def recover_the_initial_state(m, n, beta, F, Y, A):\n",
    "    time_start = time.clock()\n",
    "    flag = vector(ZZ, n)\n",
    "    for ni in range(0, n):\n",
    "        Q = Matrix(ZZ, n, n)\n",
    "        Q[0, n - 1] = (-F.list()[0]) % m\n",
    "        for i in range(1, n):\n",
    "            Q[i, i - 1] = 1\n",
    "            Q[i, n - 1] = (-F.list()[i]) % m\n",
    "        for d in range(25, 40):\n",
    "            Q_power = identity_matrix(ZZ, n)\n",
    "            recovered_a = vector(ZZ, n)\n",
    "            recovered_z = vector(ZZ, n)\n",
    "            L = Matrix(ZZ, d + 1, d + 1)\n",
    "            for i in range(1, n):\n",
    "                Q_power = (Q_power * Q) % m\n",
    "            for i in range(n, d):\n",
    "                Q_power = (Q_power * Q) % m\n",
    "                b = 0\n",
    "                for j in range(0, n):\n",
    "                    L[j + 1, i + 1] = Q_power[j, 0]\n",
    "                    b = b + Q_power[j, 0] * Y[j, ni]\n",
    "                L[0, i + 1] = (2^beta * (Y[i, ni] - b)) % m + 2^(beta - 1)\n",
    "            \n",
    "            L[0, 0] = 2^(beta - 1)\n",
    "            for i in range(1, n + 1):\n",
    "                L[0, i] = 2^(beta - 1)\n",
    "                L[i, i] = 1\n",
    "            for i in range(n + 1, d + 1):\n",
    "                L[i, i] = m\n",
    "            \n",
    "            reduced_L = L.BKZ(block_size = 20)\n",
    "            if reduced_L[0, 0] == -2^(beta - 1):\n",
    "                for i in range(0, n):\n",
    "                    recovered_z[i] = reduced_L[0, i + 1] + 2^(beta - 1)\n",
    "                    recovered_a[i] = Y[i, ni] * 2^(beta) + recovered_z[i]\n",
    "            if reduced_L[0, 0] == 2^(beta - 1):\n",
    "                for i in range(0, n):\n",
    "                    recovered_z[i] = -reduced_L[0, i + 1] + 2^(beta - 1)\n",
    "                    recovered_a[i] = Y[i, ni] * 2^(beta) + recovered_z[i]\n",
    "            \n",
    "            tol = 0\n",
    "            for i in range(0, n):\n",
    "                if recovered_a[i] == A[i, ni]:\n",
    "                    tol = tol + 1\n",
    "            if tol == n:\n",
    "                break\n",
    "                \n",
    "        flag[ni] = d\n",
    "    print(\"Success to recover the initial states by \", flag, \" truncated digits!\")\n",
    "    time_end = time.clock()\n",
    "    print('Step 4: time of recovering the initial states: ', time_end - time_start, 's\\n')\n",
    "    return flag\n",
    "\n",
    "# Recover the state transition matrix\n",
    "def recover_the_state_transition_matrix(m, n, A, C):\n",
    "    A0 = A[0:n]\n",
    "    A1 = A[1:n + 1]\n",
    "    recovered_C = (A0^(-1) * A1) % m\n",
    "    if recovered_C == C:\n",
    "        print(\"Success to recover the state transition matrix!\")\n",
    "    else:\n",
    "        print(\"Fail to recover the state transition matrix!\")\n",
    "    return recovered_C == C\n",
    "    \n",
    "    \n",
    "m = 2^31 - 1\n",
    "n = 16\n",
    "I = identity_matrix(ZZ, n)\n",
    "R.<x> = PolynomialRing(Integers(m))\n",
    "c = vector(ZZ, n)\n",
    "tol = 0\n",
    "r = 170\n",
    "t = 65\n",
    "beta = 14\n",
    "d = n * r + t - n\n",
    "\n",
    "count = vector(ZZ, r)\n",
    "sum_count = vector(ZZ, r)\n",
    "\n",
    "for i in range(0, 100):\n",
    "    time_start = time.clock()\n",
    "    print(i)\n",
    "    find_polynomial = 0\n",
    "    while find_polynomial == 0:\n",
    "        C = random_matrix(ZZ, n, n, x = 0, y = m)\n",
    "        F = x * I - C\n",
    "        find_polynomial = F.determinant().is_irreducible()\n",
    "    \n",
    "    di = r + t - n\n",
    "    A = random_matrix(ZZ, di, n, x = 0, y = m)\n",
    "    Y = Matrix(ZZ, di, n)\n",
    "    Z = Matrix(ZZ, di, n)\n",
    "    D = vector(ZZ, n)\n",
    "    for i in range(0, n):\n",
    "        D[i] = r\n",
    "    for i in range(n * r, d):\n",
    "        p = randint(1, n - 1)\n",
    "        D[p] = D[p] + 1\n",
    "        \n",
    "    for i in range(1, di):\n",
    "        A[i] = (A[i - 1] * C) % m\n",
    "    for i in range(0, n):\n",
    "        for j in range(0, D[i]):\n",
    "            Y[j, i] = A[j, i] >> beta\n",
    "            Z[j, i] = A[j, i] % (2^beta)\n",
    "    \n",
    "    count, reduced_L, flag_m = recover_the_modulus(m, n, r, t, D, A, Y)\n",
    "    flag_f = recover_the_coefficients(reduced_L, m, F.determinant(), r, t)\n",
    "    sum_count = count + sum_count\n",
    "    print(\"sum_count == \", sum_count, \"\\n\")\n",
    "    if flag_m > 0 and flag_f > 0:\n",
    "        flag_a = recover_the_initial_state(m, n, beta, F.determinant(), Y, A)\n",
    "        flag_C = recover_the_state_transition_matrix(m, n, A, C)\n",
    "    time_end = time.clock()\n",
    "    print('Total time: ', time_end - time_start, 's\\n\\n')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "SageMath 9.1",
   "language": "sage",
   "name": "sagemath"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
